import cv2
from PIL import Image, ImageDraw, ImageFont
import numpy as np
import matplotlib.pyplot as plt
import os

def process_handwriting(image_path, output_dir='chars'):
    # 创建输出目录并清空旧文件
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
    else:
        # 清空目录中已有文件
        for filename in os.listdir(output_dir):
            file_path = os.path.join(output_dir, filename)
            try:
                if os.path.isfile(file_path):
                    os.unlink(file_path)
            except Exception as e:
                print(f'无法删除文件 {file_path}: {e}')

    # 1. 读取原图并转换为灰度图（增加高斯模糊降噪）
    original = cv2.imread(image_path)
    gray = cv2.cvtColor(original, cv2.COLOR_BGR2GRAY)
   # gray = cv2.GaussianBlur(gray, (3, 3), 0)
    
    # 2. 二值化（文字白色，背景黑色）
    _, binary = cv2.threshold(gray, 120, 255, cv2.THRESH_BINARY_INV)

    # 3. 形态学处理
    # 3.1 方形腐蚀去噪点（3x3矩形内核）
    kernel_erode = cv2.getStructuringElement(cv2.MORPH_RECT, (3, 3))
    eroded = cv2.erode(binary, kernel_erode, iterations=4)  # 减少迭代次数避免过度腐蚀
    
    # 3.2 圆形膨胀恢复笔画（5x5椭圆内核）
    kernel_dilate = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (7, 7))
    dilated = cv2.dilate(eroded, kernel_dilate, iterations=5)
    
    # 3.3 中值滤波去小白点（内核调整为3x3）
    median = cv2.medianBlur(dilated, 5)
    
    # 3.4 闭运算填充内部（7x7椭圆内核）
    kernel_close = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (10, 10))
    closed = cv2.morphologyEx(median, cv2.MORPH_CLOSE, kernel_close, iterations=4)
    
    # 4. Canny边缘检测（调整阈值）
    edges = cv2.Canny(closed, 40, 120)
    
    # 5. 查找轮廓并处理
    contours, _ = cv2.findContours(edges.copy(), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    
    # 按从左到右排序轮廓（便于阅读）
    contours = sorted(contours, key=lambda c: cv2.boundingRect(c)[0])
    
    # 6. 准备结果图像
    result = gray.copy()
    result = cv2.cvtColor(result, cv2.COLOR_GRAY2BGR)
    
    # 7. 保存每个字符并绘制结果
    char_count = 0
    for i, contour in enumerate(contours):
        if cv2.contourArea(contour) > 80:  # 调整面积阈值
            x, y, w, h = cv2.boundingRect(contour)
            
            # 7.1 保存单个字符（使用原始灰度图区域）
            char_img = gray[y:y+h, x:x+w]
            # 对字符图像进行增强
            char_img = cv2.resize(char_img, (64, 64), interpolation=cv2.INTER_CUBIC)
            char_img = cv2.equalizeHist(char_img)
            cv2.imwrite(os.path.join(output_dir, f'char_{i:03d}.png'), char_img)
            char_count += 1
            
            # 7.2 在原图上绘制绿色边框
            cv2.rectangle(result, (x, y), (x+w, y+h), (0, 255, 0), 2)
    
    # 8. 添加识别结果信息
    info_text = f"Recognized {char_count} characters"
    cv2.putText(result, info_text, (10, 30), 
                cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 255, 0), 2)
    
    # 9. 显示处理过程
    plt.figure(figsize=(16, 9))
    
    # 设置中文字体
    plt.rcParams['font.sans-serif'] = ['SimHei']  # 用黑体
    plt.rcParams['axes.unicode_minus'] = False     # 解决负号显示问题
    
    steps = [gray, binary, eroded, dilated, median, closed, edges, result]
    titles = [
        '1. 原图灰度图',
        '2. 二值化（白字黑底）',
        '3. 方形腐蚀去噪',
        '4. 圆形膨胀恢复',
        '5. 中值滤波去白点',
        '6. 闭运算填充粘连',
        '7. Canny边缘检测',
        '8. 识别结果（绿色框）'
    ]
    
    for i, (img, title) in enumerate(zip(steps, titles)):
        plt.subplot(2, 4, i+1)
        if img.ndim == 2:
            plt.imshow(img, cmap='gray')
        else:
            plt.imshow(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
        plt.title(title, fontsize=10)
        plt.axis('off')
    
    plt.tight_layout()
    plt.show()
    
    # 10. 保存最终结果
    cv2.imwrite(os.path.join(output_dir, 'final_result.jpg'), result)
    return result

# 使用示例
process_handwriting('hanzi1.jpg')